#include "acl/acl.h"
#include "aclnn/acl_meta.h"
#include "torch_npu/csrc/core/npu/NPUStream.h"
#include <iostream>
#include <pybind11/pybind11.h>
#include <torch/extension.h>

namespace py = pybind11;

namespace ops
{
uint64_t getCurrentStream()
{
    return (uint64_t)(c10_npu::getCurrentNPUStream().stream());
}
} // namespace ops

struct MyAclOpExecutor
{
    aclOpExecutor *executor;
    MyAclOpExecutor() = default;
    uint64_t getPtr()
    {
        return (uint64_t)(&executor);
    }

    uint64_t get()
    {
        return (uint64_t)(executor);
    }
};

struct MyAclTensor
{
    aclTensor *ptr = nullptr;
    MyAclTensor(const at::Tensor &x)
    {
        auto sizes = x.sizes();
        auto shape_int64 = (int64_t *)(sizes.data());
        auto shape_size = sizes.size();
        auto type_meta = x.dtype();
        aclDataType acl_data_type;
        if (type_meta.name() == "c10::Half")
        {
            acl_data_type = aclDataType::ACL_FLOAT16;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "float")
        {
            acl_data_type = aclDataType::ACL_FLOAT;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "double")
        {
            acl_data_type = aclDataType::ACL_DOUBLE;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "signed char")
        {
            acl_data_type = aclDataType::ACL_INT8;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "unsigned char")
        {
            acl_data_type = aclDataType::ACL_UINT8;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "short int")
        {
            acl_data_type = aclDataType::ACL_INT16;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "int")
        {
            acl_data_type = aclDataType::ACL_INT32;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "long int")
        {
            acl_data_type = aclDataType::ACL_INT64;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else if (type_meta.name() == "bool")
        {
            acl_data_type = aclDataType::ACL_BOOL;
            std::cout << "类型" << type_meta.name() << std::endl;
        }
        else
        {
            ptr = nullptr;
            return;
        }
        // auto acl_data_type = (aclDataType)dataType;
        ptr = aclCreateTensor(shape_int64,              // view shape的维度值。
                              shape_size,               // view shape的维度个数
                              acl_data_type,            // 数据类型
                              nullptr,                  // 各维度的内存跨度数据
                              0,                        // tensor首元素相对于storage的偏移。
                              aclFormat::ACL_FORMAT_ND, // tensor内存排布格式
                              shape_int64,              // tensor的存储shape的维度值。
                              shape_size,               // tensor的存储shape的维度数。
                              x.data_ptr()              // tensor在Device侧的存储地址。
        );
    }

    ~MyAclTensor()
    {
        if (ptr)
        {
            aclDestroyTensor(ptr);
            ptr = nullptr;
        }
    }

    uint64_t get()
    {
        return (uint64_t)ptr;
    }
};

PYBIND11_MODULE(ascend_ops, m)
{
    m.doc() = "ascend ops pybind11 interfaces";
    m.def("getCurrentStream", &ops::getCurrentStream);
    py::class_<MyAclOpExecutor>(m, "AclOpExecutor")
        .def(py::init<>())
        .def("getPtr", &MyAclOpExecutor::getPtr)
        .def("get", &MyAclOpExecutor::get);
    py::class_<MyAclTensor>(m, "AclTensor")
        .def(py::init<const at::Tensor &>())
        .def("get", &MyAclTensor::get);
}
